import sys
import os
# 添加当前目录到Python路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from scipy.optimize import curve_fit
from scipy.integrate import odeint
import numpy as np
import pendulum_config

def fit_model(t, theta_data, theta0, omega0):
    """
    拟合阻尼系数k
    """
    def model(t, k):
        y0 = [theta0, omega0]
        sol = odeint(pendulum_config.pendulum_eq, y0, t, args=(k,))
        return sol[:, 0]
    
    # 使用curve_fit进行参数拟合
    popt, pcov = curve_fit(model, t, theta_data, p0=[0.1])
    k_fitted = popt[0]
    k_error = np.sqrt(np.diag(pcov))[0]
    
    return k_fitted, k_error